import torch.nn as nn

class SimpleUNet(nn.ModuleDict):
    def __init__(self,channels,timestep,imagesize):
        super().__init__()
        self.channels=channels
        self.timestep=timestep
        self.imagesize=imagesize
        # Encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(self.channels, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        
        # Middle
        self.middle = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.ReLU(),
        )
        
        # Decoder
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            
            nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, self.channels, kernel_size=3, padding=1),
        )
        
        # Time embedding
        self.time_embed = nn.Sequential(
            nn.Linear(1, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
        )
        
    def forward(self, x, t):
        # Time embedding
        t = t.float().reshape(-1, 1) / self.timestep  # normalize to [0,1]
        t_embed = self.time_embed(t)
        t_embed = t_embed.view(-1, 256, 1, 1).expand(-1, 256, int(self.imagesize/4), int(self.imagesize/4))  # adjust for the middle layer
        
        # Encoder
        x = self.encoder(x)
        
        # Middle with time embedding
        x = self.middle(x) + t_embed
        
        # Decoder
        x = self.decoder(x)
        
        return x